#include "stdafx.h"

#include <WinSock2.h>
#include <Windows.h>

#include <gtest/gtest.h>

#include "XProcess"
#include "Utils.h"

using namespace std;

void listenTest()
{
    SOCKET s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

    SOCKADDR addr;
    int addrlen = sizeof(addr);
    WSAStringToAddress("127.0.0.1:8010", AF_INET, nullptr, &addr, &addrlen);

    bind(s, &addr, addrlen);

    listen(s, 10);

    sockaddr_in inaddr;
    int inaddrlen = sizeof(sockaddr_in);
    SOCKET c = accept(s, (SOCKADDR *)&inaddr, &inaddrlen);

    char buf[1024];
    int nrecv = recv(c, buf, 1024, 0);
    if (nrecv > 0)
    {
        send(c, buf, nrecv, 0);
    }

    Sleep(1000);
    closesocket(s);
    closesocket(c);
}

void connectTest()
{
    SOCKET s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

    SOCKADDR addr;
    int addrlen = sizeof(SOCKADDR);
    WSAStringToAddress("127.0.0.1:8011", AF_INET, nullptr, &addr, &addrlen);

    Sleep(1000);
    int retv = connect(s, &addr, addrlen);
    ASSERT_EQ(retv, 0) << "Failed to connect to tester host" << endl;
    if (retv == 0)
    {
        char str[] = "<?php?><?";
        char buf[1024];
        {
            int nsent = send(s, str, sizeof(str), 0);
            int nrecv = recv(s, buf, sizeof(buf), 0);
            if (nrecv > 0)
            {
                EXPECT_STRCASEEQ(" ?php?> ?", buf);
            }
        }

        WaitForSingleObject(createProcess(L"Injector.exe -n forInject.exe -undo"), -1);

        int nsent = send(s, str, sizeof(str), 0);
        int nrecv = recv(s, buf, sizeof(buf), 0);
        if (nrecv > 0)
        {
            EXPECT_STRCASEEQ("<?php?><?", buf);
        }

        Sleep(1000);
        closesocket(s);
    }
}

TEST(Injecting, InjectDll)
{
    WSADATA wsaData;
    WSAStartup(0x201, &wsaData);
    
    createProcess(L"forInject.exe");
    createProcess(L"Injector.exe -n forInject.exe");

    listenTest();
    connectTest();
}